import unittest
from flask import Flask
from flask_testing import TestCase
import app  # 注意：这里假设你的主要应用文件名为 app.py，你需要根据实际情况调整


class TestFlaskApp(TestCase):

    def create_app(self):
        # 设置Flask应用的配置
        app.app.config['TESTING'] = True
        app.app.config['WTF_CSRF_ENABLED'] = False
        app.app.config['DEBUG'] = False
        return app.app

    def test_home_page(self):
        response = self.client.get('/')
        self.assertEqual(response.status_code, 200)
        self.assert_template_used('home.html')
        # self.assertIn(b'Home Page Content', response.data)  # 修改 'Home Page Content' 为home.html中实际存在的内容

    def test_predict_page_get(self):
        response = self.client.get('/predict')
        self.assertEqual(response.status_code, 200)
        # self.assertIn(b'Predict Page Content', response.data)  # 修改 'Predict Page Content' 为predict.html中实际存在的内容

    def test_predict_page_post(self):
        with open('test_image.png', 'rb') as img:
            img_data = {
                'file': (img, 'test_image.png')
            }
            response = self.client.post('/predict', data=img_data)
            self.assertEqual(response.status_code, 200)
            self.assert_template_used('predict.html')
            # self.assertIn(b'Unable to read the file', response.data)  # 修改检查条件为你期望的输出


if __name__ == '__main__':
    unittest.main()